# 给定一个二叉树, 找到该树中两个指定节点的最近公共祖先。
#
# 百度百科中最近公共祖先的定义为：“对于有根树 T 的两个节点 p、q，最近公共祖先表示为一个节点 x，
# 满足 x 是 p、q 的祖先且 x 的深度尽可能大（一个节点也可以是它自己的祖先）。”

from typing import List, Optional, TypeVar

from com.example.tree.tree_node import TreeNode


class Solution:

    def lowestCommonAncestor3(self, root: TreeNode, p: TreeNode, q: TreeNode) -> TreeNode:
        T = TypeVar('T', bound=TreeNode)

        parent_dict = dict[T, T]()  # 存储节点和其对应父节点的字典

        def find_parent(node: T, parent: Optional[T]) -> None:  # 定义一个辅助函数，用于递归寻找节点的父节点
            if not node:
                return
            parent_dict[node.val] = parent
            find_parent(node.left, node)
            find_parent(node.right, node)

        find_parent(root, None)  # 遍历一次树，构建节点和其对应父节点的字典

        # 从目标节点p开始，逐级向上回溯直到找到q或根节点
        p_parents = set()
        while p:
            p_parents.add(p)
            p = parent_dict.get(p.val)

        # q开始逐级向上回溯直到找到q或根节点
        while q:
            if q in p_parents:
                return q
            q = parent_dict.get(q.val)
        return None

    def lowestCommonAncestor2(self, root: TreeNode, p: TreeNode, q: TreeNode) -> TreeNode:
        """
        若 root 是 p, q 的 最近公共祖先 ，则只可能为以下情况之一：
            1. p 和 q 在 root 的子树中，且分列root 的 异侧（即分别在左、右子树中）；
            2.p = root ，且 q 在 root 的左或右子树中；
            3. q = root ，且 p 在 root 的左或右子树中；
        """

        def dfs(node: TreeNode) -> TreeNode:
            if not node or node == p or node == q:
                return node
            left_res = dfs(node.left)
            right_res = dfs(node.right)
            if not left_res and not right_res:
                return None
            if not left_res:
                return right_res
            if not right_res:
                return left_res
            return node

        return dfs(root)

    """
    找到子树里面既包含 p 也包含 q
    """

    def lowestCommonAncestor1(self, root: TreeNode, p: TreeNode, q: TreeNode) -> TreeNode:

        """
        
        dfs用于遍历树，
        返回值分别为是否包含 p 和是否包含 q
        """
        res = None

        def dfs(node: TreeNode) -> (bool, bool):
            nonlocal res  # 使用nonlocal关键字声明res为外部变量
            if not node:
                return (False, False)
            left_res = dfs(node.left)
            right_res = dfs(node.right)
            # 左子树有 p 或者右子树有 p 或者自己就是 p
            # is_contains_p = (left_res[0] or right_res[0] or node == p)
            is_contains_p = (left_res[0] or right_res[0] or node.val == p.val)
            # 左子树有 q 或者右子树有 q 或者自己就是 q
            # is_contains_q = (left_res[1] or right_res[1] or node == q)
            is_contains_q = (left_res[1] or right_res[1] or node.val == q.val)
            if is_contains_p and is_contains_q and not res:  # 第一次找到祖先时更新答案
                res = node
            return is_contains_p, is_contains_q

        dfs(root)
        return res

    def lowestCommonAncestor(self, root: TreeNode, p: TreeNode, q: TreeNode) -> TreeNode:
        return self.lowestCommonAncestor3(root, p, q)


def buildTree(preorder: List[int], inorder: List[int]) -> Optional[TreeNode]:
    inorder_map = {num: i for i, num in enumerate(inorder)}

    # 用 preorder[preorder:pre_right+1] 和 inorder[in_left:in_right+1],还原二叉树
    def build(pre_left: int, pre_right: int, in_left: int, in_right: int) -> \
            Optional[TreeNode]:
        if pre_left > pre_right:
            return None
        root = TreeNode(val=preorder[pre_left])
        root_in_inorder_index = inorder_map[root.val]  # root.val 在 inorder 中的位置
        left_children_size = root_in_inorder_index - 1 - in_left + 1
        # inorder[in_left:root_in_inorder_index] 是左子树的中序; inorder[root_in_inorder_index+1:in_right] 是右子树的中序
        # preorder[pre_left + 1:pre_left + left_children_size + 1] 是左子树的前序;
        # preorder[pre_left + left_children_size + 1+1:pre_right] 是右子树的前序
        root.left = build(pre_left + 1, pre_left + left_children_size, in_left, root_in_inorder_index - 1)
        root.right = build(pre_left + left_children_size + 1, pre_right, root_in_inorder_index + 1, in_right)
        return root

    return build(0, len(preorder) - 1, 0, len(inorder) - 1) if preorder and len(preorder) else None


if __name__ == '__main__':
    pre_order = [3, 5, 6, 2, 7, 4, 1, 0, 8]
    in_order = [6, 5, 7, 2, 4, 3, 0, 1, 8]
    p, q = TreeNode(6), TreeNode(4)
    root = buildTree(pre_order, in_order)
    res = Solution().lowestCommonAncestor(root, p, q)
    print(res.val)
